# -*- coding:utf-8 -*-

import datasets.base as input_data
import tensorflow as tf
from PIL import Image
import argparse
import sys
import random

IMAGE_SIZE = 60 * 100
LABEL_SIZE = 10


def crack_captcha(captcha_image):
    # variable in the graph for input data
    x = tf.placeholder(tf.float32, [None, IMAGE_SIZE])
    y_ = tf.placeholder(tf.float32, [None, LABEL_SIZE])

    # define the model
    W = tf.Variable(tf.zeros([IMAGE_SIZE, LABEL_SIZE]))
    b = tf.Variable(tf.zeros([LABEL_SIZE]))
    y = tf.matmul(x, W) + b

    # Define loss and optimizer
    diff = tf.nn.softmax_cross_entropy_with_logits(labels=y_, logits=y)
    cross_entropy = tf.reduce_mean(diff)
    train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)

    # forword prop
    predict = tf.argmax(y, axis=1)
    expect = tf.argmax(y_, axis=1)

    # evaluate accuracy
    correct_prediction = tf.equal(predict, expect)
    accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

    meta, train_data, test_data = input_data.load_data(
        "/Users/cheng/Desktop/DeepLearning/Project/captcha-tensorflow/images/char-1-epoch-20", flatten=True)

    saver = tf.train.Saver()
    with tf.Session() as sess:
        saver.restore(sess, "model/test.ckpt")

        print 'start testing------- we have ', train_data.images.shape[0], " samples"
        number = random.randint(0, train_data.images.shape[0])
        predict_result = sess.run(predict, feed_dict={x: [train_data.images[number]]})
        print "result is ", predict_result[0]
        img = train_data.images[number].reshape(100, 60)
        input_image = Image.fromarray(img)
        input_image.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--data_dir', type=str, default='images/char-1-epoch-20/',
                        help='Directory for storing input data')
    FLAGS, unparsed = parser.parse_known_args()
    tf.app.run(main=crack_captcha, argv=[sys.argv[0]] + unparsed)
